#include <errno.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>

#include <pcap.h>
#include <signal.h>

#include <netinet/in.h>
#include <netinet/if_ether.h>
#include <netinet/ip.h>
#include <netinet/tcp.h>
#include <netinet/udp.h>

#include <sqlite3.h>

#include <pthread.h>

#include "pseudo_pkt.h"
#include "detector1.h"

#include "logger.h"

#include "queue.h"
#include "display.h"

#define PROMISCUOUS 1
#define SIZE_OF_8021Q 4


#define MAX_LINE 50

#ifdef DEBUG_MEM
long mem_alloc;
long mem_free;
long mem_alloc_tcp;
long mem_alloc_udp;
long mem_free_irc;
#endif

char *device;     /* Network interface name (ex. eth0) */
static pcap_t *pd;
extern char *saddr;
extern int sport;

/* Queue argument */
// Layer 4 Queue
const int tcpQ_buffer_size = 10000;

// Layer 7 Queue (Common)
const int Q1_buffer_size = 1000;

// Logger Queue
const int logQ_buffer_size = 10000;

// DB
//sqlite3 *conn_malware;

/* Threads */
#define NUM_THREADS    1        /* number of thread (tcp_filter, udp_filter) */
#define NUM_THREADS_APPS   1    /* number of thread (irc_detector, dns_detector, url_detector) */



/* argument options */
int cflag;

/* ###############################################
  Consumer of tcpQ
  Producer of {ircQ, urlQ ...}
  ##################################################
*/
void *tcp_filter(void *q)
{
  // Start of TCP filter thread
  printf("Starting TCP filter thread with Queue[%x] .....\n", q);
  Queue_t *Q;
  Q = (Queue_t *)q;
  
  pseudo_pkt *pkt;
  
  while(1) {
	/* Pop packet */
	pkt = (pseudo_pkt*)queue_dequeue(Q->tcpQ);

	if (pkt->payload_size == 0) {
#ifdef DEBUG_MEM
	  mem_free++;
#endif
	  free(pkt);
	  continue;
	}

	/* filtering */
	// TCP Port 21 (telnet)
	if (pkt->dport == 21) {
#ifdef DEBUG_MEM
	  mem_free++;
#endif
	  free(pkt);
	  continue;
	}
	/* TAG List
	   Init: USER(11), MODE(12), NICK(13), USERHOST(14), JOIN(15)
	   Conn: PING(21), PONG(22)
	   Chat: PRIVMSG(31), NOTICE(32)
	*/
	if ( strncmp(pkt->payload, "PING :", 6) == 0) {
	  pkt->tag = TAG_PING;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "PONG :", 6) == 0) {
	  pkt->tag = TAG_PONG;  
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "USER ", 5) == 0 ) {
	  pkt->tag = TAG_USER;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "NICK ", 5) == 0) {
	  pkt->tag = TAG_NICK;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "USERHOST ", 9) == 0) {
	  pkt->tag = TAG_USERHOST;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "JOIN #", 6) == 0) {
	  pkt->tag = TAG_JOIN;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "MODE ", 5) == 0) {
	  pkt->tag = TAG_MODE;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "PRIVMSG ", 8) == 0) {
	  pkt->tag = TAG_PRIVMSG;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, "NOTICE ", 7) == 0) {
	  pkt->tag = TAG_NOTICE;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}
	else if ( strncmp(pkt->payload, ".u PRIVMSG ", 11) == 0) {
	  pkt->tag = TAG_PRIVMSG;
	  queue_enqueue(Q->Q1, (pseudo_pkt*)pkt);
	  continue;
	}

	// free pkt
#ifdef DEBUG_MEM
	mem_free++;
#endif
	free(pkt);
  }
  /* exit */
  pthread_exit(NULL);
}

void analysis(u_char *user, const struct pcap_pkthdr *pkthdr, const u_char *packet)
{
  unsigned short ether_type;
  struct ether_header *ehdr;

  ehdr = (struct ether_header *)packet;
  ether_type = ntohs(ehdr->ether_type);

  /* struct */
  Queue_t *q;
  q = (Queue_t *)user;

  /* struct */
  struct iphdr *iph;
  u_int8_t protocol;
  u_int16_t tot_len;
  u_char *ip_packet;
  u_char *upper;
  u_char *payload;
  
  pseudo_pkt *ppkt;
  struct tcphdr *tcphdr;
  
  /* ethertype */
  if (ether_type == ETH_P_IP)
	{
	  /* ethernet header + ip payload */
	  ip_packet = (u_char *)(packet + sizeof(struct ether_header));
	  //printf("ethernet (%d)\n", sizeof(struct ether_header));
	}
  else if (ether_type == ETH_P_8021Q)
	{
	  //printf("802.1Q (%d)\n", sizeof(struct ether_header)+SIZE_OF_8021Q);
	  ip_packet = (u_char *)(packet + sizeof(struct ether_header) + SIZE_OF_8021Q);
	}
  else
	{
	  //fprintf(stderr, "unknown ethernet header: %x\n", ether_type);
	  return;
	}

  /* decode IP Header */
  iph = (struct iphdr *)ip_packet;

#ifdef DEBUG_MEM
  mem_alloc++;
#endif
  if ( (ppkt = (pseudo_pkt *)calloc(1,sizeof(pseudo_pkt)) ) == NULL) {
	fprintf(stderr, "Calloc error of pseudo_pkt");
	exit;
  }
  // timestamp
  gettimeofday(&ppkt->timestamp, NULL);
  ppkt->saddr = iph->saddr;
  ppkt->daddr = iph->daddr;
  ppkt->protocol = iph->protocol;
  tot_len = ntohs(iph->tot_len);
  ppkt->packet_size = tot_len;

  
  upper = ip_packet + (iph->ihl * 4);
  switch(ppkt->protocol)
	{
	case IPPROTO_TCP:
	  tcphdr = (struct tcphdr *)upper;

	  /* Assign pseudo packet */
	  ppkt->sport = ntohs(tcphdr->source);
	  ppkt->dport = ntohs(tcphdr->dest);
	  payload = ip_packet + (iph->ihl * 4) + (tcphdr->doff * 4);
	  tot_len = tot_len - (iph->ihl * 4) - (tcphdr->doff * 4);
	  ppkt->payload_size = tot_len;

	  // copy payload
	  if (tot_len > SAMPLE_SIZE) tot_len = SAMPLE_SIZE;
	  memcpy(ppkt->payload, payload, tot_len);

	  /* Push ppkt into TCP Queue */
	  queue_enqueue(q->tcpQ, ppkt);
#ifdef DEBUG_MEM
	  mem_alloc_tcp++;
#endif
	  break;

	default:
#ifdef DEBUG_MEM
	  mem_alloc_udp++;
	  mem_free++;
#endif
	  free(ppkt);
	}
  return;
}
	
/* signal */
void sig_int(int sig)
{
  struct pcap_stat ps;
  int rvalue;

  /* --------- Statistics -----------*/
  rvalue = pcap_stats(pd, &ps);
  printf("-------------- Statistics ------------------\n");
  printf("Received: %d\n", ps.ps_recv);
  printf("Dropped : %d\n", ps.ps_drop);
  printf("(%d)Bye!!\n", sig);
  printf("--------------------------------------------\n");

#ifdef DEBUG_MEM
  printf("Mem alloc:%ld\n", mem_alloc);
  printf("Mem free:%ld\n", mem_free);
  printf("Mem alloc tcp:%ld\n", mem_alloc_tcp);
  printf("Mem alloc udp:%ld\n", mem_alloc_udp);
  printf("Mem free irc:%ld\n", mem_free_irc);
#endif
  
  /* need thread join */
  /* since terminate abnormally */

  /* Join threads */
  pcap_close(pd);
  exit(0);
}

void usage(void)
{
  fprintf(stdout, " Usage: \n");
  fprintf(stdout, "      -i : interface name\n");
  fprintf(stdout, "      -c : number of packet count\n");
  fprintf(stdout, "      -s : server IP address(default:localhost)\n");
  fprintf(stdout, "      -p : server Port (default:9999)\n"); 
  fprintf(stdout, "      -h : usage\n");
}

int main(int argc, char *argv[])
{
  char ebuf[PCAP_ERRBUF_SIZE];
  int c, packetcnt;
  int snaplen = 1000;    /* capture size of each packet */

  int result;
  /* db variables */

  packetcnt = -1;
  /* signal process */
  signal(SIGINT, sig_int);

  
  if (argc < 2) {
	usage();
	exit(1);
  }

  /* Options */
  while( (c = getopt(argc, argv, "i:c:s:p:h")) != -1) {
	switch(c) {
	case 'i':
	  device = optarg;
	  break;
	case 'c':
	  cflag = 1;
	  packetcnt = atoi(optarg);
	  if(packetcnt <= 0) {
		fprintf(stderr, "invalid number %s", optarg);
		fprintf(stderr, "parameter must be positive number");
		exit(1);
	  }
	  break;
	case 's':    // server IP
	  saddr = optarg;
	  break;
	case 'p':   // server Port
	  sport = atoi(optarg);
	  break;
	case 'h':
	  usage();
	  exit(1);
	}
  }

  /* init Thread & Queue */
  void *tcpQ_buffer[tcpQ_buffer_size];
  void *Q1_buffer[Q1_buffer_size];
  void *logQ_buffer[logQ_buffer_size];

  queue_t _tcpQ = QUEUE_INITIALIZER(tcpQ_buffer);
  queue_t _Q1 = QUEUE_INITIALIZER(Q1_buffer);
  queue_t _logQ = QUEUE_INITIALIZER(logQ_buffer);
  
  Queue_t *q;
  if ( (q = (Queue_t *)malloc(sizeof(Queue_t)) ) == NULL) {
	fprintf(stderr, "Cannot malloc Queue_t");
	exit;
  }

  q->tcpQ = &_tcpQ;
  q->Q1 = &_Q1;
  q->logQ = &_logQ;
  /* end of queue initialize */

  pthread_t tcp_threads[NUM_THREADS];
  int i,rc;
  /* TCP filter */

  for ( i = 0; i < NUM_THREADS; i++) {
	rc = pthread_create(&tcp_threads[i], NULL, tcp_filter, (void *)q);
	if (rc) {
	  fprintf(stderr, "ERROR: return code from pthread_create is (%d)\n", rc);
	  exit;
	}
  }


  /* detector 1 */
  pthread_t detector1_threads[NUM_THREADS_APPS];
  for ( i = 0; i < NUM_THREADS_APPS; i++) {
	rc = pthread_create(&detector1_threads[i], NULL, (void *)detector1, (void *)q);
	if (rc) {
	  fprintf(stderr, "ERROR: return code from pthread_create is (%d)\n", rc);
	  exit;
	}
  }
  
  /* logger */
  pthread_t logger_thread;
  rc = pthread_create(&logger_thread, NULL, (void *)logger, (void *)q);
  if (rc) {
	fprintf(stderr, "ERROR: return code from pthread_create is (%d)\n", rc);
	exit;
  }
  
  /* open capture interface */
  pd = pcap_open_live(device, snaplen, PROMISCUOUS, 10000, ebuf);
  if (pd == NULL) {
	perror(ebuf);
	exit(-1);
  }

  /* main processing */
  if ( (result = pcap_loop(pd, packetcnt, analysis, (void *)q)) < 0) {
	printf("result:%d\n", result);
	perror("pcap_loop error");
	exit(-1);
  }

  /* need thread join */

  for (i = 0; i < NUM_THREADS; i++) {
	pthread_join(tcp_threads[i], NULL);
  }

  for (i = 0; i < NUM_THREADS_APPS; i++) {
	pthread_join(detector1_threads[i], NULL);
  }

  pcap_close(pd);
  exit(0);
}
